import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;

import java.util.ArrayDeque;
import java.util.Deque;
import java.util.HashSet;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 自定义阻塞队列和线程池，以及Worker工作线程
 */
@Slf4j
public class TestThreadPool {

    public static void main(String[] args) {
        RejectPolicy<Runnable> rejectPolicy = (queue, task) -> {
            // 拒绝策略，五选一
            // 1、死等
            queue.put(task);
            // 2、带超时的等待
            queue.offer(task, 5, TimeUnit.SECONDS);
            // 3、调用者放弃任务执行
            log.info("放弃任务: {}", task);
            // 4、调用者抛出异常
//            throw new RuntimeException("任务失败: " + task);
            // 5、调用者自己执行任务
            task.run();
        };
        ThreadPool threadPool = new ThreadPool(2, 1, TimeUnit.SECONDS, 10, rejectPolicy);
        for (int i = 0; i < 5; i++) {
            int j = i;
            threadPool.execute(() -> {
                log.info("执行任务: 【{}】", j);
            });
        }
    }

}

/**
 * 拒绝策略 函数式接口
 * @param <T>
 */
@FunctionalInterface
interface RejectPolicy<T> {

    /**
     * 拒绝逻辑
     *
     * @param queue 装有T对象的阻塞队列
     * @param task T对象
     */
    void reject(BlockingQueue<T> queue, T task);
}

@Slf4j
class ThreadPool {

    // 任务阻塞队列
    private BlockingQueue<Runnable> taskQueue;

    // 线程集合
    private HashSet<Worker> workers = new HashSet<>();

    // 核心线程数
    private int coreSize;

    // 获取任务的超时时间
    private long timeout;
    private TimeUnit unit;

    private RejectPolicy<Runnable> rejectPolicy;

    public ThreadPool(int coreSize, long timeout, TimeUnit unit, int queueCapacity, RejectPolicy<Runnable> rejectPolicy) {
        this.coreSize = coreSize;
        this.timeout = timeout;
        this.unit = unit;
        this.taskQueue = new BlockingQueue<>(queueCapacity);
        this.rejectPolicy = rejectPolicy;
    }

    public void execute(Runnable task) {
        // 如果线程数小于核心线程数，创建新的Worker线程，否则将任务存到队列中等待线程执行
        // 需要保证workers的线程安全
        synchronized (workers) {
            if (workers.size() < coreSize) {
                Worker worker = new Worker(task);
                log.info("创建新的Worker线程: {}, {}", worker, task);
                workers.add(worker);
                worker.start();
            } else {
                // 如果核心线程数够了，尝试将任务放到阻塞队列
//                taskQueue.put(task);
                taskQueue.tryPut(rejectPolicy, task);
            }
        }
    }

    @EqualsAndHashCode(callSuper = true)
    class Worker extends Thread {
        private Runnable task;

        public Worker(Runnable task) {
            this.task = task;
        }

        @Override
        public void run() {
            // 如果任务不为空，执行任务，否则从队列中取任务执行
            // 核心线程需要一直保持核心线程不结束
//            while (task != null || (task = taskQueue.take()) != null) {
            // 非核心线程poll超时则退出
            while (task != null || (task = taskQueue.poll(timeout, unit)) != null) {
                try {
                    log.info("执行任务: {}", task);
                    task.run();
                } catch (Exception e) {
                    e.printStackTrace();
                } finally {
                    task = null;
                }
            }

            // 退出上面的while循环，说明线程执行结束了，从workers中移除
            synchronized (workers) {
                log.info("移除worker线程: {}", this);
                workers.remove(this);
            }
        }
    }
}

@Slf4j
class BlockingQueue<T> {

    // 阻塞双向队列
    private Deque<T> queue;

    // 队列操作锁
    private ReentrantLock lock = new ReentrantLock();

    // 队列已满 等待条件
    private Condition fullWaitSet = lock.newCondition();

    // 队列为空 等待条件
    private Condition emptyWaitSet = lock.newCondition();

    // 容量大小
    private int capacity;

    public BlockingQueue(int capacity) {
        this.capacity = capacity;
        this.queue = new ArrayDeque<>(capacity);
    }

    // 存任务到队尾
    public void put(T element) {
        lock.lock();
        try {
            while (queue.size() == capacity) {
                log.info("线程池满，等待放入队列: {}", element);
                fullWaitSet.await();
            }
            log.info("放入队列中: {}", element);
            queue.addLast(element);
            emptyWaitSet.signalAll();
        } catch (InterruptedException e) {
            throw new RuntimeException("队列已满 中断......");
        } finally {
            lock.unlock();
        }
    }

    // 带超时时间的 存任务到队尾
    public boolean offer(T element, long timeout, TimeUnit timeUnit) {
        lock.lock();
        try {
            long nanos = timeUnit.toNanos(timeout);
            while (queue.size() == capacity) {
                if (nanos <= 0) {
                    return false;
                }
                log.info("线程池满，等待放入队列: {}", element);
                // 等待指定纳秒后，返回的是剩余的未超时的纳秒数（等待时间-经过的时间）
                // 将剩余时间重新赋值给nanos，下次就还是等剩下的时间
                nanos = fullWaitSet.awaitNanos(nanos);
            }
            log.info("放入队列中: {}", element);
            queue.addLast(element);
            emptyWaitSet.signalAll();
            return true;
        } catch (InterruptedException e) {
            throw new RuntimeException("队列已满 中断......");
        } finally {
            lock.unlock();
        }
    }

    // 取队头的任务
    public T take() {
        lock.lock();
        try {
            while (queue.isEmpty()) {
                emptyWaitSet.await();
            }
            T t = queue.removeFirst();
            fullWaitSet.signalAll();
            return t;
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        } finally {
            lock.unlock();
        }
    }

    // 取队头的任务 带超时时间
    public T poll(long timeout, TimeUnit timeUnit) {
        lock.lock();
        try {
            // 将时间转换为纳秒
            long nanos = timeUnit.toNanos(timeout);
            while (queue.isEmpty()) {
                if (nanos <= 0) {
                    return null;
                }
                // 等待指定纳秒后，返回的是剩余的未超时的纳秒数（等待时间-经过的时间）
                // 将剩余时间重新赋值给nanos，下次就还是等剩下的时间
                nanos = emptyWaitSet.awaitNanos(nanos);
            }
            T t = queue.removeFirst();
            fullWaitSet.signalAll();
            return t;
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        } finally {
            lock.unlock();
        }
    }

    public void tryPut(RejectPolicy<T> rejectPolicy, T task) {
        lock.lock();
        try {
            // 判断队列是否已满
            if (queue.size() == capacity) {
                // 调用拒绝策略
                rejectPolicy.reject(this, task);
            } else {
                log.info("放入队列中: {}", task);
                queue.addLast(task);
                emptyWaitSet.signalAll();
            }
        } finally {
            lock.unlock();
        }
    }

    public int size() {
        lock.lock();
        try {
            return queue.size();
        } finally {
            lock.unlock();
        }
    }
}
